import socket

class ServerSocket:
    def __init__(self, port: int):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.bind(('', port))
        self.socket.listen(1)

    def __del__(self):
        print('server socket closed')
        self.socket.close()
    
    def accept(self):
        ''' Accept a client connection '''
        self.client_socket, self.client_address = self.socket.accept()
        print(f'accept: {self.client_address}')

    def receive(self, len=512):
        ''' Receive state information from the device '''
        msg = self.client_socket.recv(len).decode().split(',')
        print(f'recv state from client: {msg}')
        msg = [float(e) for e in msg]
        return msg

    def send(self, msg: str):
        ''' Send actions to the device '''
        self.client_socket.send(msg.encode())
    
    def close(self):
        ''' Close server socket '''
        self.socket.__del__()

class ClientSocket:
    def __init__(self, server_ip: str, server_port: int):
        ''' Initialize a client socket and connect to the server '''
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.socket.connect((server_ip, server_port))

    def __del__(self):
        print('client socket closed')
        self.socket.close()
    
    def receive(self, len=512):
        ''' Receive actions from the agent '''
        msg = self.socket.recv(len).decode().split(',')
        print(f'recv action from server: {msg}')
        msg = [float(e) for e in msg]
        return msg

    def send(self, msg: str):
        ''' Send states to the agent '''
        self.socket.send(msg.encode())

if __name__=="__main__":
    # 开启一个server
    server = ServerSocket(8702)
    server.accept()

    state = server.receive()
    # 训练模型...
    action = '1,1'
    server.send(action)

# Client test code
# from comm import ClientSocket

# client = ClientSocket('192.168.1.103', 8702)
# state = '1, 1, 37.5, 25'
# client.send(state)

# action = client.receive()
